/******************************************************************************
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * libkperf licensed under the Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *     http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
 * PURPOSE.
 * See the Mulan PSL v2 for more details.
 * Author:
 * Create: 2025-05-13
 * Description: Benchmark program with CPU-intensive and memory-intensive workloads
 ******************************************************************************/

#include <iostream>
#include <vector>
#include <thread>
#include <atomic>
#include <chrono>
#include <random>
#include <algorithm>
#include <iomanip>
#include <cstring>
#include <unistd.h>
#include <sched.h>
#include <sys/sysinfo.h>
#include <sys/time.h>

// Global configuration
struct BenchmarkConfig {
    int cpu_threads = 0;
    int vm_threads = 0;
    int duration_seconds = 60;
    bool verbose = false;
    bool pin_threads = false;
};

// Performance metrics
struct PerformanceMetrics {
    std::atomic<uint64_t> cpu_operations{0};
    std::atomic<uint64_t> vm_operations{0};
    std::atomic<uint64_t> cpu_cycles{0};
    std::atomic<uint64_t> vm_cycles{0};
    std::chrono::steady_clock::time_point start_time;
    std::chrono::steady_clock::time_point end_time;
};

static BenchmarkConfig config;
static PerformanceMetrics metrics;
static std::atomic<bool> stop_flag{false};

// CPU-intensive workload: Matrix multiplication with optimized algorithm
class CPUIntensiveWorkload {
private:
    static constexpr size_t MATRIX_SIZE = 512;
    std::vector<std::vector<double>> matrix_a;
    std::vector<std::vector<double>> matrix_b;
    std::vector<std::vector<double>> matrix_c;
    std::mt19937 rng;

public:
    CPUIntensiveWorkload() : rng(std::random_device{}()) {
        // Initialize matrices with random values
        matrix_a.resize(MATRIX_SIZE, std::vector<double>(MATRIX_SIZE));
        matrix_b.resize(MATRIX_SIZE, std::vector<double>(MATRIX_SIZE));
        matrix_c.resize(MATRIX_SIZE, std::vector<double>(MATRIX_SIZE));

        std::uniform_real_distribution<double> dist(-100.0, 100.0);

        for (size_t i = 0; i < MATRIX_SIZE; ++i) {
            for (size_t j = 0; j < MATRIX_SIZE; ++j) {
                matrix_a[i][j] = dist(rng);
                matrix_b[i][j] = dist(rng);
                matrix_c[i][j] = 0.0;
            }
        }
    }

    void run() {
        uint64_t operations = 0;
        uint64_t cycles = 0;

        while (!stop_flag.load()) {
            auto cycle_start = std::chrono::steady_clock::now();

            // Perform matrix multiplication with cache-friendly access pattern
            for (size_t i = 0; i < MATRIX_SIZE; ++i) {
                for (size_t k = 0; k < MATRIX_SIZE; ++k) {
                    double temp = matrix_a[i][k];
                    for (size_t j = 0; j < MATRIX_SIZE; ++j) {
                        matrix_c[i][j] += temp * matrix_b[k][j];
                    }
                }
            }

            // Add some additional CPU-intensive operations
            double sum = 0.0;
            for (size_t i = 0; i < MATRIX_SIZE; ++i) {
                for (size_t j = 0; j < MATRIX_SIZE; ++j) {
                    sum += std::sin(matrix_c[i][j]) * std::cos(matrix_c[i][j]);
                }
            }

            auto cycle_end = std::chrono::steady_clock::now();
            auto cycle_duration = std::chrono::duration_cast<std::chrono::microseconds>(cycle_end - cycle_start);

            operations++;
            cycles += cycle_duration.count();

            // Small delay to prevent overwhelming the system
            std::this_thread::sleep_for(std::chrono::microseconds(100));
        }

        metrics.cpu_operations.fetch_add(operations);
        metrics.cpu_cycles.fetch_add(cycles);
    }
};

// Memory-intensive workload: Random memory access pattern
class MemoryIntensiveWorkload {
private:
    static constexpr size_t BUFFER_SIZE = 1024 * 1024 * 64; // 64MB buffer
    std::vector<uint64_t> buffer;
    std::mt19937 rng;
    std::uniform_int_distribution<size_t> index_dist;

public:
    MemoryIntensiveWorkload() : rng(std::random_device{}()), index_dist(0, BUFFER_SIZE - 1) {
        buffer.resize(BUFFER_SIZE);

        // Initialize buffer with random values
        std::uniform_int_distribution<uint64_t> value_dist(1, 1000000);
        for (size_t i = 0; i < BUFFER_SIZE; ++i) {
            buffer[i] = value_dist(rng);
        }
    }

    void run() {
        uint64_t operations = 0;
        uint64_t cycles = 0;

        while (!stop_flag.load()) {
            auto cycle_start = std::chrono::steady_clock::now();

            // Perform random memory access pattern
            uint64_t sum = 0;
            for (int i = 0; i < 10000; ++i) {
                size_t index = index_dist(rng);
                sum += buffer[index];

                // Add some memory write operations
                if (i % 100 == 0) {
                    buffer[index] = sum;
                }
            }

            // Perform sequential memory access (cache-friendly)
            for (size_t i = 0; i < BUFFER_SIZE; i += 64) { // 64-byte cache line
                sum += buffer[i];
            }

            // Perform strided memory access (cache-unfriendly)
            for (size_t i = 0; i < BUFFER_SIZE; i += 1024) {
                sum += buffer[i];
            }

            auto cycle_end = std::chrono::steady_clock::now();
            auto cycle_duration = std::chrono::duration_cast<std::chrono::microseconds>(cycle_end - cycle_start);

            operations++;
            cycles += cycle_duration.count();

            // Small delay to prevent overwhelming the system
            std::this_thread::sleep_for(std::chrono::microseconds(100));
        }

        metrics.vm_operations.fetch_add(operations);
        metrics.vm_cycles.fetch_add(cycles);
    }
};

// Bind thread to specific CPU core
void bindThreadToCore(int thread_id, int core_id) {
    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    CPU_SET(core_id, &cpuset);

    if (sched_setaffinity(0, sizeof(cpu_set_t), &cpuset) == -1) {
        std::cerr << "Failed to bind thread " << thread_id << " to core " << core_id
                  << ": " << strerror(errno) << std::endl;
    } else if (config.verbose) {
        std::cout << "Thread " << thread_id << " bound to core " << core_id << std::endl;
    }
}

// CPU-intensive worker thread
void cpuWorker(int thread_id) {
    if (config.pin_threads) {
        bindThreadToCore(thread_id, thread_id % get_nprocs());
    }

    if (config.verbose) {
        std::cout << "CPU worker thread " << thread_id << " started" << std::endl;
    }

    CPUIntensiveWorkload workload;
    workload.run();

    if (config.verbose) {
        std::cout << "CPU worker thread " << thread_id << " finished" << std::endl;
    }
}

// Memory-intensive worker thread
void vmWorker(int thread_id) {
    if (config.pin_threads) {
        bindThreadToCore(thread_id, thread_id % get_nprocs());
    }

    if (config.verbose) {
        std::cout << "VM worker thread " << thread_id << " started" << std::endl;
    }

    MemoryIntensiveWorkload workload;
    workload.run();

    if (config.verbose) {
        std::cout << "VM worker thread " << thread_id << " finished" << std::endl;
    }
}

// Print performance metrics
void printMetrics() {
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(
        metrics.end_time - metrics.start_time).count();

    double duration_seconds = duration / 1000.0;

    uint64_t total_cpu_ops = metrics.cpu_operations.load();
    uint64_t total_vm_ops = metrics.vm_operations.load();
    uint64_t total_cpu_cycles = metrics.cpu_cycles.load();
    uint64_t total_vm_cycles = metrics.vm_cycles.load();

    std::cout << std::string(80, '=') << std::endl;
    std::cout << "Benchmark Performance Results" << std::endl;
    std::cout << std::string(80, '=') << std::endl;
    std::cout << "Duration: " << std::fixed << std::setprecision(2) << duration_seconds << " seconds" << std::endl;
    std::cout << "CPU threads: " << config.cpu_threads << std::endl;
    std::cout << "VM threads: " << config.vm_threads << std::endl;
    std::cout << std::endl;

    if (config.cpu_threads > 0) {
        double cpu_ops_per_sec = total_cpu_ops / duration_seconds;
        double cpu_avg_cycle_time = total_cpu_cycles / (double)total_cpu_ops;
        double cpu_performance_score = cpu_ops_per_sec * 1000.0; // Normalized score

        std::cout << "CPU-Intensive Workload:" << std::endl;
        std::cout << "  Total operations: " << total_cpu_ops << std::endl;
        std::cout << "  Operations/second: " << std::fixed << std::setprecision(2) << cpu_ops_per_sec << std::endl;
        std::cout << "  Average cycle time: " << std::fixed << std::setprecision(2) << cpu_avg_cycle_time << " μs" << std::endl;
        std::cout << "  Performance score: " << std::fixed << std::setprecision(2) << cpu_performance_score << std::endl;
        std::cout << std::endl;
    }

    if (config.vm_threads > 0) {
        double vm_ops_per_sec = total_vm_ops / duration_seconds;
        double vm_avg_cycle_time = total_vm_cycles / (double)total_vm_ops;
        double vm_performance_score = vm_ops_per_sec * 1000.0; // Normalized score

        std::cout << "Memory-Intensive Workload:" << std::endl;
        std::cout << "  Total operations: " << total_vm_ops << std::endl;
        std::cout << "  Operations/second: " << std::fixed << std::setprecision(2) << vm_ops_per_sec << std::endl;
        std::cout << "  Average cycle time: " << std::fixed << std::setprecision(2) << vm_avg_cycle_time << " μs" << std::endl;
        std::cout << "  Performance score: " << std::fixed << std::setprecision(2) << vm_performance_score << std::endl;
        std::cout << std::endl;
    }

    // Overall system performance
    double total_performance_score = 0.0;
    if (config.cpu_threads > 0 && config.vm_threads > 0) {
        uint64_t total_cpu_ops = metrics.cpu_operations.load();
        uint64_t total_vm_ops = metrics.vm_operations.load();
        double cpu_ops_per_sec = total_cpu_ops / duration_seconds;
        double vm_ops_per_sec = total_vm_ops / duration_seconds;
        total_performance_score = (cpu_ops_per_sec + vm_ops_per_sec) * 500.0; // Combined score
    } else if (config.cpu_threads > 0) {
        uint64_t total_cpu_ops = metrics.cpu_operations.load();
        double cpu_ops_per_sec = total_cpu_ops / duration_seconds;
        total_performance_score = cpu_ops_per_sec * 1000.0;
    } else if (config.vm_threads > 0) {
        uint64_t total_vm_ops = metrics.vm_operations.load();
        double vm_ops_per_sec = total_vm_ops / duration_seconds;
        total_performance_score = vm_ops_per_sec * 1000.0;
    }

    std::cout << "Overall Performance Score: " << std::fixed << std::setprecision(2)
              << total_performance_score << std::endl;
    std::cout << std::string(80, '=') << std::endl;
}

// Print usage information
void printUsage(const char* program_name) {
    std::cout << "Usage: " << program_name << " [OPTIONS]" << std::endl;
    std::cout << std::endl;
    std::cout << "Options:" << std::endl;
    std::cout << "  --cpu N          Number of CPU-intensive threads (default: 0)" << std::endl;
    std::cout << "  --vm N           Number of memory-intensive threads (default: 0)" << std::endl;
    std::cout << "  --time N         Duration in seconds (default: 60)" << std::endl;
    std::cout << "  --pin            Pin threads to CPU cores" << std::endl;
    std::cout << "  --verbose        Enable verbose output" << std::endl;
    std::cout << "  --help           Show this help message" << std::endl;
    std::cout << std::endl;
    std::cout << "Examples:" << std::endl;
    std::cout << "  " << program_name << " --cpu 4 --vm 4 --time 30" << std::endl;
    std::cout << "  " << program_name << " --cpu 8 --time 60 --pin --verbose" << std::endl;
    std::cout << "  " << program_name << " --vm 16 --time 120" << std::endl;
    std::cout << std::endl;
    std::cout << "Workloads:" << std::endl;
    std::cout << "  CPU-intensive: Matrix multiplication with trigonometric operations" << std::endl;
    std::cout << "  Memory-intensive: Random and strided memory access patterns" << std::endl;
}

// Parse command line arguments
bool parseArguments(int argc, char* argv[]) {
    for (int i = 1; i < argc; ++i) {
        std::string arg = argv[i];

        if (arg == "--help" || arg == "-h") {
            printUsage(argv[0]);
            return false;
        } else if (arg == "--cpu") {
            if (i + 1 < argc) {
                config.cpu_threads = std::stoi(argv[++i]);
                if (config.cpu_threads < 0) {
                    std::cerr << "Error: CPU thread count must be non-negative" << std::endl;
                    return false;
                }
            } else {
                std::cerr << "Error: --cpu requires a number" << std::endl;
                return false;
            }
        } else if (arg == "--vm") {
            if (i + 1 < argc) {
                config.vm_threads = std::stoi(argv[++i]);
                if (config.vm_threads < 0) {
                    std::cerr << "Error: VM thread count must be non-negative" << std::endl;
                    return false;
                }
            } else {
                std::cerr << "Error: --vm requires a number" << std::endl;
                return false;
            }
        } else if (arg == "--time") {
            if (i + 1 < argc) {
                config.duration_seconds = std::stoi(argv[++i]);
                if (config.duration_seconds <= 0) {
                    std::cerr << "Error: Duration must be positive" << std::endl;
                    return false;
                }
            } else {
                std::cerr << "Error: --time requires a number" << std::endl;
                return false;
            }
        } else if (arg == "--pin") {
            config.pin_threads = true;
        } else if (arg == "--verbose") {
            config.verbose = true;
        } else {
            std::cerr << "Error: Unknown option '" << arg << "'" << std::endl;
            printUsage(argv[0]);
            return false;
        }
    }

    if (config.cpu_threads == 0 && config.vm_threads == 0) {
        std::cerr << "Error: At least one of --cpu or --vm must be specified" << std::endl;
        printUsage(argv[0]);
        return false;
    }

    return true;
}

int main(int argc, char* argv[]) {
    std::cout << "Benchmark Program - CPU and Memory Intensive Workloads" << std::endl;
    std::cout << "=====================================================" << std::endl;

    if (!parseArguments(argc, argv)) {
        return EXIT_FAILURE;
    }

    // Print configuration
    std::cout << "Configuration:" << std::endl;
    std::cout << "  CPU threads: " << config.cpu_threads << std::endl;
    std::cout << "  VM threads: " << config.vm_threads << std::endl;
    std::cout << "  Duration: " << config.duration_seconds << " seconds" << std::endl;
    std::cout << "  Thread pinning: " << (config.pin_threads ? "enabled" : "disabled") << std::endl;
    std::cout << "  Available CPU cores: " << get_nprocs() << std::endl;
    std::cout << std::endl;

    // Start benchmark
    std::cout << "Starting benchmark..." << std::endl;
    metrics.start_time = std::chrono::steady_clock::now();

    std::vector<std::thread> threads;

    // Start CPU-intensive threads
    for (int i = 0; i < config.cpu_threads; ++i) {
        threads.emplace_back(cpuWorker, i);
    }

    // Start memory-intensive threads
    for (int i = 0; i < config.vm_threads; ++i) {
        threads.emplace_back(vmWorker, config.cpu_threads + i);
    }

    // Wait for specified duration
    std::cout << "Running for " << config.duration_seconds << " seconds..." << std::endl;
    std::this_thread::sleep_for(std::chrono::seconds(config.duration_seconds));

    // Stop all threads
    stop_flag.store(true);

    // Wait for all threads to finish
    for (auto& thread : threads) {
        thread.join();
    }

    metrics.end_time = std::chrono::steady_clock::now();

    // Print results
    printMetrics();

    std::cout << "Benchmark completed successfully." << std::endl;
    return EXIT_SUCCESS;
}